import numpy as np
from MyMLTools.metrics import r2_score

class LinearRegression:
    def __init__(self):
        self.coef_ = None
        self.interception = None
        self._theta = None

    def fit_normal(self, X_train, y_train):
        assert X_train.shape[0] == y_train.shape[0], \
            '训练集特征矩阵（X_train）的规模必须与训练集标签向量（y_train）的规模相等'
        X_b = np.hstack([np.ones((len(X_train), 1)), X_train])
        self._theta = np.linalg.inv(X_b.T.dot(X_b)).dot(X_b.T).dot(y_train)
        self.interception = self._theta[0]
        self.coef_ = self._theta[1:]
        return self

    def predict(self, X_predict):
        assert self.interception is not None and self.coef_ is not None, \
            '先拟合，再预测！'
        assert X_predict.shape[1] == len(self.coef_), \
            '预测集 X_predict 的特征数量必须与训练集 X_train （的特征数量）相等'
        X_b = np.hstack([np.ones((len(X_predict), 1)), X_predict])
        return X_b.dot(self._theta)

    def score(self, X_test, y_test):
        y_predict = self.predict(X_test)
        return r2_score(y_test, y_predict)